#!/usr/bin/env python
# --*-- coding:utf-8 --*--
# author:g-y-b time:2020/6/1

import cv2
import numpy as np


def hough(img, threshold=300):
    thetas = np.deg2rad(np.arange(0, 180, 2))
    row, cols = img.shape
    diag_len = np.ceil(np.sqrt(row ** 2 + cols ** 2))
    rhos = np.linspace(-diag_len, diag_len, int(2 * diag_len))
    cos_t = np.cos(thetas)
    sin_t = np.sin(thetas)
    num_theta = len(thetas)
    # 构造计数矩阵
    vote = np.zeros((int(2 * diag_len), num_theta), dtype=np.uint64)
    y_inx, x_inx = np.nonzero(img)
    # 计数
    for i in range(len(x_inx)):
        x = x_inx[i]
        y = y_inx[i]
        for j in range(num_theta):
            rho = round(x * cos_t[j] + y * sin_t[j]) + diag_len
            vote[int(rho), j] += 1
    # 拿到计数大于阈值的theta和rho
    indeies = np.argwhere(vote > threshold)
    rhos_idx = indeies[:, 0]
    theta_idx = indeies[:, 1]
    return vote, rhos[rhos_idx], thetas[theta_idx]


img = cv2.imread('book.jpg')
cv2.imshow('1src', img)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
edges = cv2.Canny(gray, 70, 140)
cv2.imshow('2edge', edges)

vote, rhos, thetas = hough(edges)
vote = np.uint8(vote.T)
cv2.imshow('3vote', vote)

for rho, theta in zip(rhos, thetas):
    x_center = img.shape[1]/2
    x1 = int(x_center+500)
    x2 = int(x_center-500)
    a = np.sin(theta)
    b = np.cos(theta)
    if a == 0.0:
        a = 1e-5
    if b == 0.0:
        b = 1e-5
    y1 = int((rho-x1*b)/a)
    y2 = int((rho-x2*b)/a)
    cv2.line(img, (x1, y1), (x2, y2), (0, 255, 0), 2)

cv2.imshow('4hough', img)
print("计算结束")
cv2.waitKey(0)
